package com.zyx.flinkdemo.sql.func;

import com.zyx.flinkdemo.pojo.WaterSensor;
import com.zyx.flinkdemo.pojo.WeightedAvgAccum;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.functions.AggregateFunction;

import static org.apache.flink.table.api.Expressions.$;
import static org.apache.flink.table.api.Expressions.call;

/**
 * @author zyx
 * @since 2021/5/30 16:59
 * desc: UDTF(聚合函数)案例
 */
public class UdtfDemo {
    public static void main(String[] args) throws Exception {

        // 1.获取执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        // 2.读取端口数据并转换为JavaBean
        SingleOutputStreamOperator<WaterSensor> waterSensorDs = env
                .socketTextStream("linux201", 7777)
                .map(line -> {
                    String[] split = line.split(",");
                    return new WaterSensor(split[0],
                            Long.parseLong(split[1]),
                            Integer.parseInt(split[2]));
                });

        // 3.将流转换为动态表
        Table table = tableEnv.fromDataStream(waterSensorDs);

        // 4.先注册再使用
        tableEnv.createTemporarySystemFunction("myavg", WeightedAvg.class);
        tableEnv.createTemporaryView("water_sensor", table);

        //TableAPI
        /*table.groupBy($("id"))
                .select($("id"), call("myavg", $("vc")))
                .execute()
                .print();*/

        //SQL
        tableEnv
                .sqlQuery("select id,myavg(vc) from water_sensor group by id")
                .execute()
                .print();

        //6.执行任务
        env.execute();
    }

    public static class WeightedAvg extends AggregateFunction<Double, WeightedAvgAccum> {
        @Override
        public WeightedAvgAccum createAccumulator() {
            return new WeightedAvgAccum();
        }

        @Override
        public Double getValue(WeightedAvgAccum acc) {
            return acc.getSum() * 1D / acc.getCount();
        }

        public void accumulate(WeightedAvgAccum acc, Integer vc) {
            acc.setCount(acc.getCount() + 1);
            acc.setSum(acc.getSum() + vc);
        }
    }
}
